package org.xxd.kafka.clients.common;


import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author: XiaoDong.Xie
 * @create: 2020-09-27 11:26
 * @description:
 */
public class Cluster {
    /** 所有的broker节点 */
    private final List<Node> nodes;
    /** topic 对应所有的分区 */
    private final Map<String, List<PartitionInfo>> partitionsByTopic;
    /** topic 对应所有可用的分区 */
    private final Map<String, List<PartitionInfo>> availablePartitionsByTopic;
    /** nodeId 对应所有的Node */
    private final Map<Integer, Node> nodesById;
    /** topic 一个具体的分区 对应的PartitionInfo 信息*/
    private final Map<TopicPartition, PartitionInfo> partitionsByTopicPartition;
    /** 一个Node 有哪些分区 */
    private final Map<Integer, List<PartitionInfo>> partitionsByNode;

    public Cluster(Collection<Node> nodes,
                    Collection<PartitionInfo> partitions) {

        // make a randomized, unmodifiable copy of the nodes
        List<Node> copy = new ArrayList<>(nodes);
        Collections.shuffle(copy);
        this.nodes = Collections.unmodifiableList(copy);

        this.nodesById = new HashMap<>();
        for (Node node : nodes) {
            this.nodesById.put(node.id(), node);
        }

        // index the partitions by topic/partition for quick lookup
        this.partitionsByTopicPartition = new HashMap<>(partitions.size());
        for (PartitionInfo p : partitions) {
            this.partitionsByTopicPartition.put(new TopicPartition(p.topic(), p.partition()), p);
        }

        // index the partitions by topic and node respectively, and make the lists
        // unmodifiable so we can hand them out in user-facing apis without risk
        // of the client modifying the contents
        HashMap<String, List<PartitionInfo>> partsForTopic = new HashMap<>();
        HashMap<Integer, List<PartitionInfo>> partsForNode = new HashMap<>();
        for (Node n : this.nodes) {
            partsForNode.put(n.id(), new ArrayList<PartitionInfo>());
        }

        for (PartitionInfo p : partitions) {
            if (!partsForTopic.containsKey(p.topic())) {
                partsForTopic.put(p.topic(), new ArrayList<PartitionInfo>());
            }
            List<PartitionInfo> psTopic = partsForTopic.get(p.topic());
            psTopic.add(p);

            if (p.leader() != null) {
                List<PartitionInfo> psNode = partsForNode.get(p.leader().id());
                psNode.add(p);
            }
        }

        this.partitionsByTopic = new HashMap<>(partsForTopic.size());
        this.availablePartitionsByTopic = new HashMap<>(partsForTopic.size());

        for (Map.Entry<String, List<PartitionInfo>> entry : partsForTopic.entrySet()) {
            String topic = entry.getKey();
            List<PartitionInfo> partitionList = entry.getValue();
            this.partitionsByTopic.put(topic, Collections.unmodifiableList(partitionList));
            List<PartitionInfo> availablePartitions = new ArrayList<>();
            for (PartitionInfo part : partitionList) {
                if (part.leader() != null) {
                    availablePartitions.add(part);
                }
            }
            this.availablePartitionsByTopic.put(topic, Collections.unmodifiableList(availablePartitions));
        }

        this.partitionsByNode = new HashMap<>(partsForNode.size());
        for (Map.Entry<Integer, List<PartitionInfo>> entry : partsForNode.entrySet()) {
            this.partitionsByNode.put(entry.getKey(), Collections.unmodifiableList(entry.getValue()));
        }
    }

    public static Cluster empty() {
        return new Cluster(new ArrayList<Node>(0), new ArrayList<PartitionInfo>(0));
    }

    public static Cluster bootstrap(List<InetSocketAddress> addresses) {
        List<Node> nodes = new ArrayList<>();
        int nodeId = -1;
        for (InetSocketAddress address : addresses)
            nodes.add(new Node(nodeId--, address.getHostString(), address.getPort()));
        return new Cluster(nodes, new ArrayList<PartitionInfo>(0));
    }

    /**
     * 根据topic获取所有的分区
     * @param topic
     * @return
     */
    public List<PartitionInfo> partitionsForTopic(String topic) {
        return this.partitionsByTopic.get(topic);
    }


    /**
     * 获取topic 可用的分区
     * @param topic
     * @return
     */
    public List<PartitionInfo> availablePartitionForTopic(String topic) {
        return this.availablePartitionsByTopic.get(topic);
    }

    public List<Node> nodes() {
        return this.nodes;
    }

    public Node leaderFor(TopicPartition topicPartition) {
        PartitionInfo info = partitionsByTopicPartition.get(topicPartition);
        if (info == null)
            return null;
        else
            return info.leader();
    }
}
